
#pragma once
#ifndef UNETV1_H_
#define UNETV1_H_
#include <stdlib.h>
#include <iostream>
#include "stream_tools.h"

constexpr uint32_t N_ICH = 4;
constexpr uint32_t N_OCH = 6;
constexpr uint32_t N_IH = 5;
constexpr uint32_t N_IW = 5;
constexpr uint32_t K = 3;
constexpr uint32_t P = 1;
constexpr uint32_t S = 2;
constexpr uint32_t O_P = 1;
constexpr uint32_t N_OH = (N_IH - 1) * S + K - 2 * P + O_P;
constexpr uint32_t N_OW = (N_IW - 1) * S + K - 2 * P + O_P;
constexpr uint32_t VEC_LEN = N_OH * N_OW;
constexpr uint32_t P_ICH = 2;
constexpr uint32_t P_OCH = 3;
constexpr uint32_t FOLD_I = N_ICH / P_ICH;
constexpr uint32_t FOLD_O = N_OCH / P_OCH;
static_assert(N_ICH % P_ICH == 0, "N_ICH must be divisible by P_ICH");
static_assert(N_OCH % P_OCH == 0, "N_OCH must be divisible by P_OCH");
constexpr uint32_t N_IM = N_IH * N_IW;
constexpr uint32_t N_OM = N_OH * N_OW;
constexpr unsigned BIT_ACTV = 32; 
constexpr unsigned BIT_WGHT = 8; 
constexpr unsigned BIT_CONV = 32;
extern ap_uint<P_OCH * P_ICH * BIT_WGHT> weight_reshape[FOLD_O][FOLD_I][K*K];

constexpr unsigned CONV1_ICH = 4;
constexpr unsigned CONV1_OCH = 8;
constexpr unsigned CONV1_H = 64;
constexpr unsigned CONV1_W = 64;
constexpr unsigned CONV1_P_ICH = 2;
constexpr unsigned CONV1_P_OCH = 4;
extern ap_uint<CONV1_P_OCH * CONV1_P_ICH * BIT_WGHT> conv1_weight[CONV1_OCH/CONV1_P_OCH][CONV1_ICH/CONV1_P_ICH][K*K];

constexpr unsigned POOL1_CH = 8;
constexpr unsigned POOL1_IH = 64;
constexpr unsigned POOL1_IW = 64;
constexpr unsigned POOL1_OH = 32;
constexpr unsigned POOL1_OW = 32;
constexpr unsigned POOL1_P_CH = 4;

constexpr unsigned CONV2_ICH = 8;
constexpr unsigned CONV2_OCH = 16;
constexpr unsigned CONV2_H = 32;
constexpr unsigned CONV2_W = 32;
constexpr unsigned CONV2_P_ICH = 4;
constexpr unsigned CONV2_P_OCH = 4;
extern ap_uint<CONV2_P_OCH * CONV2_P_ICH * BIT_WGHT> conv2_weight[CONV2_OCH/CONV2_P_OCH][CONV2_ICH/CONV2_P_ICH][K*K];

constexpr unsigned POOL2_CH = 16;
constexpr unsigned POOL2_H = 16;
constexpr unsigned POOL2_W = 16;
constexpr unsigned POOL2_P_CH = 4;

constexpr unsigned CONV3_ICH = 16;
constexpr unsigned CONV3_OCH = 32;
constexpr unsigned CONV3_H = 16;
constexpr unsigned CONV3_W = 16;
constexpr unsigned CONV3_P_ICH = 4;
constexpr unsigned CONV3_P_OCH = 4;
extern ap_uint<CONV3_P_OCH * CONV3_P_ICH * BIT_WGHT> conv3_weight[CONV3_OCH/CONV3_P_OCH][CONV3_ICH/CONV3_P_ICH][K*K];

constexpr unsigned POOL3_CH = 32;
constexpr unsigned POOL3_H = 8;
constexpr unsigned POOL3_W = 8;
constexpr unsigned POOL3_P_CH = 4;

constexpr unsigned CONV4_ICH = 32;
constexpr unsigned CONV4_OCH = 64;
constexpr unsigned CONV4_H = 8;
constexpr unsigned CONV4_W = 8;
constexpr unsigned CONV4_P_ICH = 4;
constexpr unsigned CONV4_P_OCH = 4;
extern ap_uint<CONV4_P_OCH * CONV4_P_ICH * BIT_WGHT> conv4_weight[CONV4_OCH/CONV4_P_OCH][CONV4_ICH/CONV4_P_ICH][K*K];

constexpr unsigned POOL4_CH = 64;
constexpr unsigned POOL4_H = 4;
constexpr unsigned POOL4_W = 4;
constexpr unsigned POOL4_P_CH = 4;

constexpr unsigned CENTER_ICH = 64;
constexpr unsigned CENTER_OCH = 128;
constexpr unsigned CENTER_H = 4;
constexpr unsigned CENTER_W = 4;
constexpr unsigned CENTER_P_ICH = 4;
constexpr unsigned CENTER_P_OCH = 4;
extern ap_uint<CENTER_P_OCH * CENTER_P_ICH * BIT_WGHT> center_weight[CENTER_OCH/CENTER_P_OCH][CENTER_ICH/CENTER_P_ICH][K*K];

constexpr unsigned DECONV1_ICH = 128;
constexpr unsigned DECONV1_OCH = 64;
constexpr unsigned DECONV1_IH = 4;
constexpr unsigned DECONV1_IW = 4;
constexpr unsigned DECONV1_OH = 8;
constexpr unsigned DECONV1_OW = 8;
constexpr unsigned DECONV1_P_ICH = 4;
constexpr unsigned DECONV1_P_OCH = 4;
extern ap_uint<DECONV1_P_OCH * DECONV1_P_ICH * BIT_WGHT> deconv1_weight[DECONV1_OCH/DECONV1_P_OCH][DECONV1_ICH/DECONV1_P_ICH][K*K];

constexpr unsigned CONV5_ICH = 128;
constexpr unsigned CONV5_OCH = 64;
constexpr unsigned CONV5_H = 8;
constexpr unsigned CONV5_W = 8;
constexpr unsigned CONV5_P_ICH = 4;
constexpr unsigned CONV5_P_OCH = 4;
extern ap_uint<CONV5_P_OCH * CONV5_P_ICH * BIT_WGHT> conv5_weight[CONV5_OCH/CONV5_P_OCH][CONV5_ICH/CONV5_P_ICH][K*K];

constexpr unsigned DECONV2_ICH = 64;
constexpr unsigned DECONV2_OCH = 32;
constexpr unsigned DECONV2_IH = 8;
constexpr unsigned DECONV2_IW = 8;
constexpr unsigned DECONV2_OH = 16;
constexpr unsigned DECONV2_OW = 16;
constexpr unsigned DECONV2_P_ICH = 4;
constexpr unsigned DECONV2_P_OCH = 4;
extern ap_uint<DECONV2_P_OCH * DECONV2_P_ICH * BIT_WGHT> deconv2_weight[DECONV2_OCH/DECONV2_P_OCH][DECONV2_ICH/DECONV2_P_ICH][K*K];

constexpr unsigned CONV6_ICH = 64;
constexpr unsigned CONV6_OCH = 32;
constexpr unsigned CONV6_H = 16;
constexpr unsigned CONV6_W = 16;
constexpr unsigned CONV6_P_ICH = 4;
constexpr unsigned CONV6_P_OCH = 4;
extern ap_uint<CONV6_P_OCH * CONV6_P_ICH * BIT_WGHT> conv6_weight[CONV6_OCH/CONV6_P_OCH][CONV6_ICH/CONV6_P_ICH][K*K];

constexpr unsigned DECONV3_ICH = 32;
constexpr unsigned DECONV3_OCH = 16;
constexpr unsigned DECONV3_IH = 16;
constexpr unsigned DECONV3_IW = 16;
constexpr unsigned DECONV3_OH = 32;
constexpr unsigned DECONV3_OW = 32;
constexpr unsigned DECONV3_P_ICH = 4;
constexpr unsigned DECONV3_P_OCH = 4;
extern ap_uint<DECONV3_P_OCH * DECONV3_P_ICH * BIT_WGHT> deconv3_weight[DECONV3_OCH/DECONV3_P_OCH][DECONV3_ICH/DECONV3_P_ICH][K*K];

constexpr unsigned CONV7_ICH = 32;
constexpr unsigned CONV7_OCH = 16;
constexpr unsigned CONV7_H = 16;
constexpr unsigned CONV7_W = 16;
constexpr unsigned CONV7_P_ICH = 4;
constexpr unsigned CONV7_P_OCH = 4;
extern ap_uint<CONV7_P_OCH * CONV7_P_ICH * BIT_WGHT> conv7_weight[CONV7_OCH/CONV7_P_OCH][CONV7_ICH/CONV7_P_ICH][K*K];

constexpr unsigned DECONV4_ICH = 16;
constexpr unsigned DECONV4_OCH = 8;
constexpr unsigned DECONV4_IH = 32;
constexpr unsigned DECONV4_IW = 32;
constexpr unsigned DECONV4_OH = 64;
constexpr unsigned DECONV4_OW = 64;
constexpr unsigned DECONV4_P_ICH = 4;
constexpr unsigned DECONV4_P_OCH = 4;
extern ap_uint<DECONV4_P_OCH * DECONV4_P_ICH * BIT_WGHT> deconv4_weight[DECONV4_OCH/DECONV4_P_OCH][DECONV4_ICH/DECONV4_P_ICH][K*K];

constexpr unsigned CONV8_ICH = 16;
constexpr unsigned CONV8_OCH = 8;
constexpr unsigned CONV8_H = 32;
constexpr unsigned CONV8_W = 32;
constexpr unsigned CONV8_P_ICH = 4;
constexpr unsigned CONV8_P_OCH = 4;
extern ap_uint<CONV8_P_OCH * CONV8_P_ICH * BIT_WGHT> conv8_weight[CONV8_OCH/CONV8_P_OCH][CONV8_ICH/CONV8_P_ICH][K*K];


// Function declarations
void unetv1(data_stream<BIT_ACTV>& in, axiu_stream<BIT_CONV>& out);
void unetv1_golden(data_stream<BIT_ACTV>& in, data_stream<BIT_CONV>& out, 
                   const ap_int<BIT_WGHT> weight[N_OCH][K*K][N_ICH]);

#endif